from idaapi import *
from idc import *

def write_data_to_file(data,filename,filecount=1):
    realfilename=filename+".txt"
    f = open(realfilename,"ab+")
    fenge = "======="+ str(filecount) +"========"
    f.write(fenge)
    f.write(data)
    f.close()

if __name__=="__main__":
    filepath = "E:/works/baiduditu/data/"
    MD5KeyBeforeFile = filepath+"MD5KeyBefore"
    MD5KeyAfterFile = filepath+"MD5KeyAfter"
    CompressBeforeFile =filepath+"CompressBefore"
    CompressAfterFile = filepath+"CompressAfter"
    BeforeEAFile = filepath+"BeforeEA"
    AfterEAFile = filepath+"AfterEA"
    beforexorFile = filepath+"beforexor"

    bpMD5Key = GetBptEA(0)
    beforxor = GetBptEA(1)
    bpEA = GetBptEA(2)
    bpCompress = GetBptEA(3)

    print "ResumeProcess before"

    ResumeProcess()
    print "ResumeProcess after"

    while 1:
        ret = GetDebuggerEvent(WFNE_SUSP, -1)
        print("BREAKPOINT reached")
        if(ret == BREAKPOINT):
            pc = GetRegValue("PC")

            if(pc == bpMD5Key):
                afterMD5Data = GetRegValue("R1")
                beforeMD5Data = GetRegValue("R2")
                beforeMD5DataLen = GetRegValue("R3")
                print("MD5 reached")
                strbeforeMD5Data = DbgRead(beforeMD5Data,beforeMD5DataLen)
                write_data_to_file(strbeforeMD5Data,MD5KeyBeforeFile,1)
                StepOver()
                ret = GetDebuggerEvent(WFNE_SUSP, -1)
                if 1:
                    strafterMD5Data = DbgRead(afterMD5Data,32)
                    write_data_to_file(strafterMD5Data, MD5KeyAfterFile, 1)

            if(pc==bpCompress):
                print("bpCompress reached")
                input = GetRegValue("R0")
                inputLen = GetRegValue("R1")
                output = GetRegValue("R2")
                strinput = DbgRead(input,inputLen)
                write_data_to_file(strinput, CompressBeforeFile, 1)
                StepOver()
                ret = GetDebuggerEvent(WFNE_SUSP, -1)
                if 1:
                    outputlen = GetRegValue("R3")
                    print (str(outputlen))
                    stroutput = DbgRead(output,outputlen)
                    write_data_to_file(stroutput,CompressAfterFile)

            if(pc==beforxor):
                print("beforxor reached")
                beforxordata = GetRegValue("R5")
                beforxordatalen = GetRegValue("R1")
                strbeforxordata=DbgRead(beforxordata,beforxordatalen)
                write_data_to_file(strbeforxordata,beforexorFile)

            if(pc == bpEA):
                print("bpEA reached")
                EABeforeData = GetRegValue("R1")
                EABeforeDataLen = GetRegValue("R2")
                EAAfterData = GetRegValue("R0")
                strEABeforeData=DbgRead(EABeforeData, EABeforeDataLen)
                write_data_to_file(strEABeforeData, BeforeEAFile)
                StepOver()
                ret = GetDebuggerEvent(WFNE_SUSP, -1)
                if 1:
                    print(str(1))
                    strEAAfterData = DbgRead(EAAfterData, len(strEABeforeData)*2)
                    write_data_to_file(strEAAfterData,AfterEAFile)
        ResumeProcess()


